import numpy as np
import networkx as nx
import logging
import os


from project_setting import DATA_PATH
from knowgraph.node2vec import Node2Vec
from knowgraph.generator_walks import GeneratorWalks
from utils.neo4j_query_utils import update_info_to_neo4j,run_neo4j_query,get_neo4j_dirver_instance
from common.random_walk_query import NEORJ_QUERY_update_cluster_to_enterprise,NEO4J_ENTERPRISER_BIDDING_RELATION



from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score

logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)








def kmeans_cluster(embeddings, k=-1, clusterProp='clusterId'):
    """
    Given a list of dicts like {"nodeId" 1, "embedding": [1.0, 0.1, ...]},
    generate a list of dicts like {"nodeId": 1, "valueMap": {"clusterId": 2}}
    """
    print("Performing K-Means clustering (k={}, clusterProp='{}')"
          .format(k, clusterProp))
    X = np.array([e["embedding"] for e in embeddings])

    if k == -1:
        max_sil_score = 0
        best_cluster_num = -1
        for cluster_num in range(10,200):
            kmeans = KMeans(n_clusters=cluster_num).fit(X)
            score = silhouette_score(X, kmeans.labels_)
            if max_sil_score <= score:
                max_sil_score = score
                best_cluster_num = cluster_num
            print("当前轮廓系数:%s ,最大的轮廓系数:%s" % (str(score),str(max_sil_score)))
        k = best_cluster_num
    print("最好的聚类个数："+ str(k))

    kmeans = KMeans(n_clusters=k).fit(X)
    results = []
    for idx, cluster in enumerate(kmeans.predict(X)):
        results.append({ "nodeId": embeddings[idx]["nodeId"],
                         "valueMap": { clusterProp: int(cluster) }})
    print("...clustering completed.")

    from sklearn.manifold import TSNE
    np.set_printoptions(suppress=True)
    model = TSNE()
    Y = model.fit_transform(X)  # 将X降维(默认二维)后保存到Y中

    '''可视化过程'''
    from matplotlib import pyplot as plt
    plt.scatter(Y[:, 0], Y[:, 1], 20, kmeans.labels_)  # labels为每一行对应标签，20为标记大小
    plt.savefig("cluster.png")  # 保存图片
    plt.show()
    return results


if __name__ == "__main__":
    save_base_path = os.path.join(DATA_PATH,'node2vec')

    # 先导出所有的投标企业关联关系
    enterprise_bidding_relation_save_file = os.path.join(save_base_path,'enterprise_bidding/enterprise_bidding_relation.txt')

    if os.path.exists(enterprise_bidding_relation_save_file):
        train_data = []
        with open(enterprise_bidding_relation_save_file,'r',encoding='utf-8') as fread:
            for line in fread:
                train_data.append(line.strip())
    else:
        os.makedirs(os.path.dirname(enterprise_bidding_relation_save_file),exist_ok=True)
        train_data = run_neo4j_query(NEO4J_ENTERPRISER_BIDDING_RELATION, get_neo4j_dirver_instance())
        train_data = [" ".join([str(j) for j in i]) for i in train_data]
        duplicate_train_data = set()
        for data in train_data:
            if data not in duplicate_train_data:
                duplicate_train_data.add(data)
        train_data = list(duplicate_train_data)

        with open(enterprise_bidding_relation_save_file,'w',encoding='utf-8') as fwriter:
            for item in train_data:
                fwriter.write(item + "\n")



    random_walk_param = {
        "walk_length":10,
        "num_walks":200000,
        'p':1,
        'q':1,
    }

    G = nx.parse_edgelist(train_data,
                          create_using=nx.DiGraph(), data=[('weight', int)], nodetype=None)
    # 判断是否有随机游走的文件，没有则从新生成
    random_walk_save_path = os.path.join(save_base_path,'enterprise_bidding/')
    random_walks_obj = GeneratorWalks(random_walk_save_path,**random_walk_param)
    random_walks = random_walks_obj.load_walks(G)

    model = Node2Vec(word2vec_save_basepath=random_walk_save_path)
    model.train(sentences=random_walks)
    embedding = model.get_embeddings()


    cluster_embeding = []
    for key,value in embedding.items():
        cluster_item = {}
        cluster_item['nodeId'] = key
        cluster_item['embedding'] = value
        cluster_embeding.append(cluster_item)

    clusters = kmeans_cluster(cluster_embeding,k=25)
    update_info_to_neo4j(get_neo4j_dirver_instance(), NEORJ_QUERY_update_cluster_to_enterprise,{"updates":clusters})
